use axum::{
    body::Body as AxumBody,
    extract::{Request, State},
    http::{HeaderMap, StatusCode},
    response::Response,
    routing::any,
    Router,
};
use clap::Parser;
use dotenvy::dotenv;
use futures::StreamExt;
use reqwest::Client;
use std::{env, net::SocketAddr, sync::Arc};
use tower_http::trace::TraceLayer;
use tracing::{debug, info};

#[derive(Parser, Debug)]
#[command(author, version, about)]
struct Args {
    /// 启用调试模式
    #[arg(short, long)]
    debug: bool,
}

#[derive(Clone)]
struct AppState {
    api_key: String,
    ollama_base_url: String,
    client: Client,
    debug: bool,
}

#[tokio::main]
async fn main() {
    // 解析命令行参数
    let args = Args::parse();

    // 初始化环境变量
    dotenv().ok();

    // 初始化日志
    if args.debug {
        std::env::set_var("RUST_LOG", "debug");
    } else {
        std::env::set_var("RUST_LOG", "info");
    }
    tracing_subscriber::fmt::init();

    // 获取配置
    let api_key = env::var("API_KEY").expect("API_KEY must be set");
    let ollama_base_url = env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://localhost:11434".to_string());
    let port = env::var("PORT").unwrap_or_else(|_| "3500".to_string()).parse::<u16>().expect("PORT must be a valid number");

    // 创建共享状态
    let state = Arc::new(AppState {
        api_key,
        ollama_base_url,
        client: Client::new(),
        debug: args.debug,
    });

    // 创建路由
    let app = Router::new()
        .route("/{path}", any(proxy_handler))
        .fallback(proxy_handler)
        .layer(TraceLayer::new_for_http())
        .with_state(state);

    // 启动服务器
    let addr = SocketAddr::from(([0, 0, 0, 0], port));
    info!("Server listening on {}", addr);
    if args.debug {
        debug!("Debug mode enabled");
    }

    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn proxy_handler(
    State(state): State<Arc<AppState>>,
    req: Request,
) -> Result<Response, StatusCode> {
    if state.debug {
        debug!("收到请求: {} {}", req.method(), req.uri());
        debug!("请求头: {:?}", req.headers());
    }

    // 验证 API Key
    let auth_header = req
        .headers()
        .get("Authorization")
        .and_then(|h| h.to_str().ok())
        .ok_or(StatusCode::UNAUTHORIZED)?;

    let expected_auth = format!("Bearer {}", state.api_key);
    if auth_header != expected_auth {
        if state.debug {
            debug!("认证失败: 无效的 API Key");
        }
        return Err(StatusCode::UNAUTHORIZED);
    }

    // 构建转发请求的URL
    let path = req.uri().path();
    let query = req.uri().query().unwrap_or("");
    let url = if query.is_empty() {
        format!("{}{}", state.ollama_base_url, path)
    } else {
        format!("{}{}?{}", state.ollama_base_url, path, query)
    };

    if state.debug {
        debug!("转发请求到: {}", url);
    }

    // 提取原始请求的headers和method
    let method = req.method().clone();
    let mut headers = HeaderMap::new();
    for (key, value) in req.headers() {
        if key != "host" {  // 跳过host header
            headers.insert(key.clone(), value.clone());
        }
    }

    // 读取请求体
    let body_bytes = axum::body::to_bytes(req.into_body(), usize::MAX)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    if state.debug {
        if let Ok(body_str) = String::from_utf8(body_bytes.clone().to_vec()) {
            debug!("请求体: {}", body_str);
        }
    }

    // 创建新的请求
    let client_req = state
        .client
        .request(method, &url)
        .headers(headers)
        .body(body_bytes)
        .send()
        .await
        .map_err(|_| StatusCode::BAD_GATEWAY)?;

    if state.debug {
        debug!("收到上游响应: {} {:?}", client_req.status(), client_req.headers());
    }

    // 构建响应
    let status = client_req.status();
    let headers = client_req.headers().clone();

    // 将响应流式传输回客户端
    let stream = client_req.bytes_stream().map(move |result| {
        if state.debug {
            if let Ok(ref bytes) = result {
                if let Ok(text) = String::from_utf8(bytes.to_vec()) {
                    debug!("响应数据片段: {}", text);
                }
            }
        }
        result.map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "stream error"))
    });
    let body = AxumBody::from_stream(stream);

    let mut response = Response::new(body);
    *response.status_mut() = status;
    *response.headers_mut() = headers;

    Ok(response)
}
